package edu.nepu.flink.api.state;

import edu.nepu.flink.api.bean.WaterSensor;
import org.apache.flink.api.common.functions.FlatMapFunction;
import org.apache.flink.api.common.functions.RichMapFunction;
import org.apache.flink.api.common.state.*;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.BroadcastStream;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.co.BroadcastProcessFunction;
import org.apache.flink.util.Collector;

import java.util.Iterator;

/**
 * @Date 2024/3/2 10:42
 * @Created by chenshuaijun
 */
public class OperatorBroadcastStateDemo {

    public static void main(String[] args) throws Exception {

        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        env.setParallelism(1);
        /**
         * TODO 现在我们来使用动态划定水位线的功能来体会广播状态的使用
         * 我们要解决的问题是，主流的数据正常的输入，另外开一条流作为配置流，通过配置流来告诉主流，当水位值超过多少的时候产生报警
         */

        DataStreamSource<String> mainSource = env.socketTextStream("hadoop102", 9999);
        DataStreamSource<String> waringSource = env.socketTextStream("hadoop102", 8888);
        SingleOutputStreamOperator<WaterSensor> mapSource = mainSource.flatMap(new FlatMapFunction<String, WaterSensor>() {
            @Override
            public void flatMap(String value, Collector<WaterSensor> out) throws Exception {
                try {
                    String[] split = value.split(",");
                    if (split.length == 3) {
                        out.collect(new WaterSensor(split[0], Long.parseLong(split[1]), Integer.parseInt(split[2])));
                    }
                } catch (Exception e) {
                    // 如果因为数据输入错误，产生的异常直接忽略掉
                }
            }
        });
        MapStateDescriptor<String, Integer> mapStateDescriptor = new MapStateDescriptor<>("waring-threshold", Types.STRING, Types.INT);
        BroadcastStream<String> broadcast = waringSource.broadcast(mapStateDescriptor);
        mapSource.connect(broadcast).process(new MyBroadcastThreshold(mapStateDescriptor)).print();
        env.execute();
    }

    static class MyBroadcastThreshold extends BroadcastProcessFunction<WaterSensor, String, String>{
        MapStateDescriptor<String, Integer> mapStateDescriptor;

        public MyBroadcastThreshold(MapStateDescriptor<String, Integer> mapStateDescriptor) {
            this.mapStateDescriptor = mapStateDescriptor;
        }

        @Override
        public void processElement(WaterSensor value, BroadcastProcessFunction<WaterSensor, String, String>.ReadOnlyContext ctx, Collector<String> out) throws Exception {
            ReadOnlyBroadcastState<String, Integer> broadcastState = ctx.getBroadcastState(mapStateDescriptor);
            Integer threshold = broadcastState.get("threshold");
            threshold = threshold == null? 0 : threshold;
            if (value.vc > threshold){
                out.collect(value+"的vc大于当前threshold"+threshold);
            }
        }

        @Override
        public void processBroadcastElement(String value, BroadcastProcessFunction<WaterSensor, String, String>.Context ctx, Collector<String> out) throws Exception {
            BroadcastState<String, Integer> broadcastState = ctx.getBroadcastState(mapStateDescriptor);
            broadcastState.put("threshold",Integer.parseInt(value));
        }

    }

}
